// Copyright (c) 2025 Huawei Technologies Co., Ltd.
// openUBMC is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.
//
// Description: 实现KVM客户端websocket连接

#ifndef WEBSOCKET_CLIENT_HPP
#define WEBSOCKET_CLIENT_HPP

#include <boost/beast/core.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/websocket/ssl.hpp>
#include <boost/asio/ssl.hpp>
#include <boost/asio/strand.hpp>
#include <cstdlib>
#include <iostream>
#include <thread>
#include <memory>
#include <string>
#include <queue>
#include "HeartbeatManager.hpp"

namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace net = boost::asio;
namespace ssl = boost::asio::ssl;
using tcp = boost::asio::ip::tcp;

constexpr int MAX_READ_BUFFER = 16 * 1024 * 1024;
constexpr int MAX_WRITE_BUFFER = 1024 * 1024;

class WebSocketClient {
public:
    WebSocketClient();
    virtual ~WebSocketClient();

    void Connect(std::string host, std::string port, std::string uri);
    void Disconnect();
    void Send(const std::vector<uint8_t>& binary_data);
    void Run();
    void StartHeartbeat(std::chrono::seconds interval);
    void StopHeartbeat();

protected:
    virtual void OnMessage(const std::string& message) = 0;
    virtual void OnClose(const websocket::close_reason& cr) = 0;
    virtual void OnError(const std::exception& error) = 0;

private:
    net::io_context io_context_;
    ssl::context ssl_context_;
    tcp::resolver resolver_;
    boost::optional<websocket::stream<ssl::stream<beast::tcp_stream>>> ws_;
    beast::multi_buffer buffer_;
    std::string host_;
    std::string port_;
    std::string uri_;
    HeartbeatManager heartbeat_manager_;

    // 消息队列相关成员
    bool isWriting_ = false; // 标记是否有正在进行的写操作
    std::queue<std::vector<uint8_t>> sendQueue_;
    std::mutex sendMutex_;

    void ProcessSendQueue(); // 处理发送队列

    void SendHeartbeat(const std::vector<uint8_t>& message);

    void on_resolve(const boost::system::error_code& ec, tcp::resolver::results_type results);
    void on_connect(const boost::system::error_code& ec, const tcp::endpoint& ep);
    void on_ssl_handshake(const boost::system::error_code& ec);
    void on_handshake(const boost::system::error_code& ec);
    void do_read();
    void on_read(const boost::system::error_code& ec, std::size_t bytes_transferred);
    void on_write(beast::error_code ec, std::size_t bytes_transferred);
    void on_close(const websocket::close_reason& cr);
    void on_error(const std::exception& ec);
};

#endif  // WEBSOCKET_CLIENT_HPP